# -*- coding: utf-8 -*-
"""
文件名称: data_filter.py

功能概述:
    启发于 sub_exp_1 系列实验结论, 在保持所有诊断题 (diagnosis / quality+diagnosis) 的前提下,
    对冗余度较高的其他题目类型 (judge / compare / choice) 进行比例过滤, 以减少数据规模、
    保留关键信息并减小训练开销。

    具体思路:
        1. 利用 support_cleaned JSON 文件中每张图像对应的疾病标签(label)来统计“疾病组合”的出现频次,
           以此衡量其稀有度 (出现次数越少 => 越稀有)。
           - support_cleaned 格式示例:
                {
                  "aptos_001639a390f0.png": {
                    "image_path": "/path/to/aptos_001639a390f0.png",
                    "label": ["增殖型糖尿病视网膜病变"],
                    "quality": null
                  },
                  "messidor_20060411_58550_0200_PP.png": {
                    "image_path": ".../messidor_20060411_58550_0200_PP.png",
                    "label": null,
                    "quality": "图像质量不可评估，无法用于眼底疾病诊断。"
                  },
                  ...
                }
           - 我们将 label 列表按字典序排序后用 "_" join 成一个字符串, 作为“疾病 key”,
             统计该 key 在所有图像中的出现次数。
           - 对于 label 为 None 或空列表的样本:
                * 不参与疾病稀有度统计 (即不生成疾病 key、也不计频次)
                * 但记录仍存在于支持数据中, 表示图像在其它任务中可能可用
                * 若后续“错误地”参与 FILTER_TYPES 题型的过滤过程, 将在过滤阶段触发报错
                  (因为其不具备稀有度/keep_ratio, 属于流程设计错误或数据对齐问题)

        2. 按疾病驱动的分段线性策略 (默认模式, 未指定 fixed_ratio 时启用)

           基于“疾病 key”的绝对频次进行分段线性策略, 为每个疾病 key 分配保留概率 keep_ratio,
           再映射到图像级别:
           - 记某疾病组合的频次为 f, 所有疾病按频次从高到低排序:
                 (d1, f1), (d2, f2), (d3, f3), (d4, f4), (d5, f5), ...
           - Top3 疾病 d1, d2, d3:
               * 无论频次多少, 一律 keep_ratio = 0.05 (强力压缩最常见的 3 类)
           - 其他疾病按频次分三段:
               ① 低频区: f ≤ 20
                    keep_ratio = 1.0  (全部保留, 不动长尾)
               ② 中频区: 20 < f ≤ 1000
                    线性从 1.0 衰减到 0.2:
                        t = (f - 20) / (1000 - 20)
                        keep(f) = 1.0 - t * (1.0 - 0.2) = 1.0 - 0.8 * t
               ③ 高频区: f > 1000
                    以“第五大频次 f5”作为高频段的上界:
                        - 若 f5 > 1000:
                              t = (f - 1000) / (f5 - 1000)
                              keep(f) = 0.2 - t * (0.2 - 0.05)
                                      = 0.2 - 0.15 * t
                              并与 0.05 取最大值:
                                  keep(f) = max(0.05, keep(f))
                        - 若 f5 ≤ 1000 (极端情况, 小数据集):
                              对非 Top3 且 f>1000 的疾病统一 keep_ratio = 0.2 兜底

           - 最终保证:
                * 所有 keep_ratio ∈ (0, 1]
                * Top3 疾病固定为 0.05
                * 低频疾病大部分接近 1.0, 随着频次增加平滑衰减到 0.2
                * 头部类进一步衰减到 0.05, 但不过度“削掉”尾部

           - 当 fixed_ratio 为 None 时, 本脚本在输出目录(out_dir=输出文件所在目录)额外写出一个
             JSON 日志文件, 记录:
                {
                  "等效删除比": <全局等效删除比例>,
                  "疾病情况": [
                    [疾病名称, 频次, 删除比例],
                    ...
                  ]
                }
             其中删除比例 = 1 - keep_ratio, 列表按疾病频次从高到低排序。

        3. 固定比例模式 (指定 --fixed-ratio 时启用)

           - 若命令行指定 --fixed-ratio=<p>, 且 0 < p ≤ 1.0,
             则完全跳过疾病频次与稀有度逻辑, 使用一个全局统一的保留比例:
                 image_keep_ratio[img_name] = p
             其中 img_name 为 support_cleaned 中的每一个 key (即所有支持集中出现过的图像,
             无论其 label 为 None 还是 list, 都会被赋予相同的 keep_ratio)。
           - 这样可以:
               * 做对照实验: “按疾病自适应抽样” vs “统一 p% 抽样”
               * 保证 FILTER_TYPES 中涉及到的所有图像, 只要在 support_cleaned 中出现过, 就都有 keep_ratio,
                 不再依赖疾病标签是否为空。
           - 在固定比例模式下, 不输出“疾病情况”的 per-disease 删除比例日志, 仅使用统一比例。

        4. 对 Alpaca 格式 JSON 数据进行过滤:
           - 输入 JSON 顶层为 list[dict], 每个元素为一条样本, 格式示例:
                [
                  {
                    "question_type": "compare",
                    "instruction": "...",
                    "input": "",
                    "output": "...",
                    "images": ["/path/to/1979327_l_eye.jpg"]
                  },
                  ...
                ]
           - 通过 images[0] 取出图像路径, 用 os.path.basename 提取文件名,
             再在 image_keep_ratio 中查找保留概率 keep_ratio。
           - 对以下题目类型进行随机过滤 (全局变量 FILTER_TYPES):
                FILTER_TYPES = ["judge", "compare", "choice"]
             * 若 question_type 不在 FILTER_TYPES 中 (如 diagnosis / quality+diagnosis / quality),
               则一律保留, 不做过滤。
             * 若 question_type 在 FILTER_TYPES 中, 则:
                 - 生成 [0,1) 的随机数 r
                 - 若 r ≤ keep_ratio: 保留该样本
                 - 否则: 丢弃该样本 (不写入输出列表)
             * 若某图像在 FILTER_TYPES 中出现, 但 img_name 不在 image_keep_ratio 中,
               则认为流程存在错误 (支持文件缺失或病种统计逻辑不一致), 在过滤阶段直接 raise。

        5. 节流与等效删除比例:
           在完成整份数据的过滤后, 额外计算并打印“等效删除比例”, 方便直观评估本次过滤强度:
               - 全局等效删除比例:
                     global_delete_ratio = filtered / total
                 表示在所有样本中, 实际被删除的比例。
               - 可过滤题型子集上的等效删除比例:
                     filtered_sub / total_sub  (仅统计 question_type ∈ FILTER_TYPES 的部分)
                 表示在真正参与随机裁剪的题目子集中, 实际删除比例有多大。

使用示例:
    # 自适应疾病频率模式 (+ 输出疾病级别删除比例日志)
    python data_filter.py \
        --input-jsonl  /path/to/alpaca_raw.json \
        --output-jsonl /path/to/alpaca_filtered.json \
        --support-cleaned-path /path/to/support_cleaned.json \
        --seed 42

    # 固定比例模式 (例如统一保留 20%, 不输出 per-disease 删除日志)
    python data_filter.py \
        --input-jsonl  /path/to/alpaca_raw.json \
        --output-jsonl /path/to/alpaca_filtered_p20.json \
        --support-cleaned-path /path/to/support_cleaned.json \
        --fixed-ratio 0.2 \
        --seed 42
"""

import argparse
import json
import logging
import os
import random
from collections import Counter
from pathlib import Path
from typing import Any, Dict, List, Tuple


# -----------------------------
# 全局配置
# -----------------------------

# 需要按稀有度进行过滤的题目类型
FILTER_TYPES = ["judge", "compare", "choice"]

# 分段线性规则参数 (与 disease_keep_ratio_stats.py 保持一致)

# 低频保留阈值: f <= LOW_FULL_KEEP_THRESHOLD -> keep_ratio = 1.0
LOW_FULL_KEEP_THRESHOLD = 20

# 中频区间: (MID_START, MID_END] 做线性插值
MID_START = 20
MID_END = 1000

# Top3 疾病固定保留比例
TOP3_KEEP_RATIO = 0.05

# 中频段末端的目标保留比例
MID_END_KEEP_RATIO = 0.20

# 高频段最低保留比例 (Top3 之外的疾病不会低于该值)
HIGH_FREQ_MIN_KEEP_RATIO = 0.05

logger = logging.getLogger(__name__)
if not logger.handlers:
    _handler = logging.StreamHandler()
    _formatter = logging.Formatter(
        fmt="[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    _handler.setFormatter(_formatter)
    logger.addHandler(_handler)
    logger.setLevel(logging.INFO)


# -----------------------------
# 支持信息读取与稀有度统计
# -----------------------------

def load_support_cleaned(support_path: Path) -> Dict[str, Any]:
    """
    读取支持 JSON 文件, 返回原始 dict。

    参数:
        support_path: 支持 JSON 文件路径

    返回:
        support_dict: 原始 JSON 解析后的 dict

    异常:
        - 文件不存在或无法读取时抛出异常
        - JSON 解析失败时抛出异常
        - 顶层结构不是 dict 时抛出异常
    """
    if not support_path.is_file():
        raise FileNotFoundError(f"support_cleaned_path 文件不存在: {support_path}")

    try:
        text = support_path.read_text(encoding="utf-8")
        data = json.loads(text)
    except Exception as e:
        raise RuntimeError(f"读取或解析支持 JSON 失败: path={support_path}, error={e}") from e

    if not isinstance(data, dict):
        raise RuntimeError(
            f"支持 JSON 顶层结构必须为 dict, 实际为: {type(data)} (path={support_path})"
        )

    return data


def build_disease_key_and_freq(
    support_dict: Dict[str, Any]
) -> Tuple[Dict[str, str], Dict[str, int]]:
    """
    基于支持 JSON, 统计 “疾病 key” 的出现频次, 并为每个图像文件名构建对应疾病 key 映射。

    疾病 key 的构造方式:
        - 对于 support_dict 中的每个条目:
            * key: 图像文件名, 如 "aptos_001639a390f0.png"
            * value: 记录 dict, 其中 "label" 字段为疾病标签列表
        - 对 label 列表执行:
            * 若 label 为 None 或为空列表: 认为该样本没有可用于稀有度统计的疾病信息,
              仅记录在支持集中存在, 但不生成疾病 key / 频次 (即跳过统计);
              这类图像后续不会出现在 image_to_disease_key 中, 因而也不会有 keep_ratio,
              若在过滤阶段对其进行 FILTER_TYPES 题型的过滤尝试, 将触发报错。
            * 若 label 为非空 list: 对 label 做 sorted, 然后用 "_" join, 得到 disease_key 字符串。

    返回:
        image_to_disease_key: {image_filename: disease_key}
            - 仅包含那些 label 非空、可参与稀有度统计的图像
        disease_freq:         {disease_key: 出现次数}

    异常:
        - 任一条记录的 label 字段既不是 None 也不是 list 时 raise
    """
    image_to_disease_key: Dict[str, str] = {}
    freq_counter: Counter = Counter()

    for img_name, info in support_dict.items():
        if not isinstance(info, dict):
            raise RuntimeError(
                f"support_cleaned 中 key={img_name!r} 对应的 value 非 dict 类型: {type(info)}"
            )

        labels = info.get("label", None)

        # 情况 1: label 为 None 或空列表 -> 不参与稀有度统计, 也不进入 image_to_disease_key
        if labels is None:
            logger.debug(
                "support_cleaned: key=%s 的 label 为 None, 跳过稀有度统计",
                img_name,
            )
            continue
        if isinstance(labels, list) and len(labels) == 0:
            logger.debug(
                "support_cleaned: key=%s 的 label 为空列表, 跳过稀有度统计",
                img_name,
            )
            continue

        # 情况 2: 正常的非空 list
        if not isinstance(labels, list):
            raise RuntimeError(
                f"support_cleaned 中 key={img_name!r} 的 label 类型异常, "
                f"期望为 list / None, 实际为: {type(labels)}, value={labels!r}"
            )

        labels_sorted = sorted(str(x) for x in labels)
        disease_key = "_".join(labels_sorted)

        image_to_disease_key[img_name] = disease_key
        freq_counter[disease_key] += 1

    disease_freq: Dict[str, int] = dict(freq_counter)
    if not disease_freq:
        raise RuntimeError("在 support_cleaned 中未统计到任何疾病 key, 请检查输入文件与 label 字段")

    logger.info(
        "从支持文件中统计到 %d 个不同疾病 key, 总计 %d 条图像记录(具有非空标签)",
        len(disease_freq),
        len(image_to_disease_key),
    )
    return image_to_disease_key, disease_freq


def build_disease_keep_ratio(
    disease_freq: Dict[str, int],
) -> Dict[str, float]:
    """
    根据疾病在数据集中的“绝对频次”动态分配 keep_ratio, 更温和地削减高频疾病,
    避免对头部类进行过度删除, 同时尽量不动长尾。

    设计 (频次记为 f, 分三段 + Top3 特殊处理):

        1) 排序:
            - 按频次从高到低排序得到列表:
                  sorted_desc = [(d1, f1), (d2, f2), (d3, f3), (d4, f4), (d5, f5), ...]

            - Top3 疾病 d1, d2, d3:
                  keep_ratio = TOP3_KEEP_RATIO (默认 0.05)

            - f5 作为高频段插值的上界:
                  若疾病数 >= 5, 则 f5 = 第五大频次
                  否则退化为 max_freq (防御性处理)

        2) 分段规则 (针对非 Top3 疾病):

            (a) 低频区: f <= LOW_FULL_KEEP_THRESHOLD (默认 20)
                -> keep_ratio = 1.0

            (b) 中频区: LOW_FULL_KEEP_THRESHOLD < f <= MID_END (默认 20 < f <= 1000)
                -> 从 1.0 线性衰减到 MID_END_KEEP_RATIO (默认 0.2):

                   t = (f - MID_START) / (MID_END - MID_START)
                   keep(f) = 1.0 - t * (1.0 - MID_END_KEEP_RATIO)

            (c) 高频区: f > MID_END (默认 f > 1000)

                - 若 f5 > MID_END:
                    t = (f - MID_END) / (f5 - MID_END)
                    keep(f) = MID_END_KEEP_RATIO - t * (MID_END_KEEP_RATIO - HIGH_FREQ_MIN_KEEP_RATIO)
                             = 0.2 - t * 0.15
                    最后与 HIGH_FREQ_MIN_KEEP_RATIO 取最大值:
                        keep(f) = max(HIGH_FREQ_MIN_KEEP_RATIO, keep(f))

                - 若 f5 <= MID_END (极端情况):
                    -> 非 Top3 的 f>1000 疾病一律 keep_ratio = MID_END_KEEP_RATIO 兜底。

    返回:
        disease_keep_ratio: {disease_key: keep_ratio ∈ (0, 1]}
    """
    if not disease_freq:
        raise RuntimeError("build_disease_keep_ratio: disease_freq 为空, 无法分配 keep_ratio")

    # 按频次从高到低排序
    sorted_desc = sorted(disease_freq.items(), key=lambda x: x[1], reverse=True)
    max_freq = sorted_desc[0][1]
    min_freq = min(disease_freq.values())

    # Top3 疾病
    top3_keys = {k for k, _ in sorted_desc[:3]}

    # 第五大频次 (若不足 5 个疾病, 退化为 max_freq)
    if len(sorted_desc) >= 5:
        f5 = float(sorted_desc[4][1])
    else:
        f5 = float(max_freq)

    logger.info(
        "疾病频次范围: min_freq=%d, max_freq=%d, 第五大频次 f5=%.0f",
        min_freq,
        max_freq,
        f5,
    )
    logger.info(
        "Top3 疾病固定 keep_ratio=%.3f, 分别为: %s",
        TOP3_KEEP_RATIO,
        ", ".join(list(top3_keys)),
    )

    disease_keep_ratio: Dict[str, float] = {}

    for d_key, freq in disease_freq.items():
        f = float(freq)

        # 1) Top3 疾病: 固定 0.05
        if d_key in top3_keys:
            keep = TOP3_KEEP_RATIO

        # 2) 低频区: f <= 20 -> 全保留
        elif f <= LOW_FULL_KEEP_THRESHOLD:
            keep = 1.0

        # 3) 中频区: (20, 1000] -> 线性从 1.0 衰减到 0.2
        elif f <= MID_END:
            if MID_END == MID_START:
                keep = MID_END_KEEP_RATIO
            else:
                t = (f - MID_START) / (MID_END - MID_START)
                keep = 1.0 - t * (1.0 - MID_END_KEEP_RATIO)

        # 4) 高频区: f > 1000
        else:
            if f5 > MID_END:
                denom = (f5 - MID_END) if f5 != MID_END else 1.0
                t = (f - MID_END) / denom
                # 防御: t 理论上在 [0,1], 这里裁剪一下避免数值抖动
                if t < 0.0:
                    t = 0.0
                elif t > 1.0:
                    t = 1.0
                keep = MID_END_KEEP_RATIO - t * (MID_END_KEEP_RATIO - HIGH_FREQ_MIN_KEEP_RATIO)
                if keep < HIGH_FREQ_MIN_KEEP_RATIO:
                    keep = HIGH_FREQ_MIN_KEEP_RATIO
            else:
                # 极端情况: 没有真正意义上的高频段, 统一用 0.2 兜底
                keep = MID_END_KEEP_RATIO

        # 数值安全: keep ∈ (0,1]
        if keep > 1.0:
            keep = 1.0
        if keep <= 0.0:
            keep = HIGH_FREQ_MIN_KEEP_RATIO

        disease_keep_ratio[d_key] = float(keep)

    # 简单日志: 看看 keep_ratio 大致分布 (按疾病数)
    summary_counter: Dict[float, int] = {}
    for v in disease_keep_ratio.values():
        summary_counter[v] = summary_counter.get(v, 0) + 1

    logger.info("频率驱动 keep_ratio 分布概况(按疾病数):")
    for kr, cnt in sorted(summary_counter.items(), key=lambda x: x[0], reverse=True):
        logger.info("  keep_ratio=%.3f -> disease_count=%d", kr, cnt)

    return disease_keep_ratio


def build_image_keep_ratio(
    image_to_disease_key: Dict[str, str],
    disease_keep_ratio: Dict[str, float],
) -> Dict[str, float]:
    """
    将疾病级别的 keep_ratio 映射到图像级别。

    参数:
        image_to_disease_key: {image_filename: disease_key}
            - 只包含那些具有非空 label(可统计稀有度)的图像
        disease_keep_ratio:  {disease_key: keep_ratio}

    返回:
        image_keep_ratio:    {image_filename: keep_ratio}
    """
    image_keep_ratio: Dict[str, float] = {}

    for img_name, d_key in image_to_disease_key.items():
        if d_key not in disease_keep_ratio:
            raise RuntimeError(
                f"疾病 key={d_key!r} 在 disease_keep_ratio 中未找到, "
                f"请检查构建逻辑 (img_name={img_name!r})"
            )
        image_keep_ratio[img_name] = disease_keep_ratio[d_key]

    logger.info(
        "已为 %d 张图像分配 keep_ratio (基于疾病稀有度, 仅含有 label 的样本)",
        len(image_keep_ratio),
    )
    return image_keep_ratio


def build_image_keep_ratio_fixed(
    support_dict: Dict[str, Any],
    fixed_ratio: float,
) -> Dict[str, float]:
    """
    在“固定比例模式”下, 为支持文件中出现的所有图像分配统一的 keep_ratio。

    参数:
        support_dict: 支持 JSON (key 为图像文件名, value 为信息 dict, 不强制要求含 label)
        fixed_ratio:  固定保留比例, 要求 0 < fixed_ratio <= 1.0

    返回:
        image_keep_ratio: {image_filename: fixed_ratio}

    说明:
        - 不区分是否有 label; 只要图像出现在支持文件中, 就会被赋予同样的保留概率。
        - 适合做对照实验: “按疾病自适应抽样” vs “统一 p% 抽样”。
    """
    if not (0.0 < fixed_ratio <= 1.0):
        raise ValueError(
            f"fixed_ratio 必须在 (0,1] 区间内, 当前值为: {fixed_ratio!r}"
        )

    image_keep_ratio: Dict[str, float] = {}
    for img_name, info in support_dict.items():
        if not isinstance(info, dict):
            raise RuntimeError(
                f"support_cleaned 中 key={img_name!r} 对应的 value 非 dict 类型: {type(info)}"
            )
        image_keep_ratio[img_name] = float(fixed_ratio)

    logger.info(
        "固定比例模式启用: 为 %d 张图像统一分配 keep_ratio=%.3f (不依赖疾病稀有度)",
        len(image_keep_ratio),
        fixed_ratio,
    )
    return image_keep_ratio


# -----------------------------
# Alpaca JSON 过滤主逻辑（list[dict]）
# -----------------------------

def filter_alpaca_json(
    input_path: Path,
    output_path: Path,
    image_keep_ratio: Dict[str, float],
    seed: int = 42,
) -> Dict[str, Any]:
    """
    读取 Alpaca JSON 文件(顶层 list[dict]), 根据题目类型与图像对应疾病稀有度 / 固定比例进行过滤,
    输出新的 JSON 文件(同样是 list[dict])。

    过滤策略:
        - 若 record["question_type"] 不在 FILTER_TYPES 中, 则一律保留。
          (例如 diagnosis / quality+diagnosis / quality 等题型始终不过滤)
        - 若 question_type 在 FILTER_TYPES 中, 则:
            * 通过 record["images"][0] 获取图像路径, 使用 os.path.basename 得到文件名 img_name
            * 在 image_keep_ratio 中查找 keep_ratio:
                - 若找不到, 说明:
                    1) 该图像在支持文件中缺失, 或
                    2) 频率模式下该图像 label 为 None/空、未参与统计,
                       但仍被用于 FILTER_TYPES 中 (设计错误)。
                  此时直接 raise RuntimeError, 提醒检查数据与过滤逻辑。
            * 生成随机数 r ~ U[0,1).
                - 若 r <= keep_ratio: 保留该条样本 (写入输出列表)
                - 否则: 丢弃

    参数:
        input_path:   输入 Alpaca JSON 文件路径, 顶层 list[dict]
        output_path:  输出过滤后的 JSON 文件路径, 顶层 list[dict]
        image_keep_ratio: 图像级别的保留概率映射 (可来自自适应模式或固定比例模式)
        seed:         随机种子, 默认 42

    返回:
        stats: 一个字典, 包含整体统计信息, 例如:
            {
              "total": int,
              "kept": int,
              "filtered": int,
              "global_delete_ratio": float,
              "filtered_types_delete_ratio": float,
              "type_stats": {
                qtype: {"total": int, "kept": int, "filtered": int}, ...
              }
            }

    额外统计 (节流相关):
        - total / kept / filtered: 全局样本数、保留数、删除数。
        - type_stats: 按 question_type 统计 total / kept / filtered。
        - global_delete_ratio: 全局等效删除比例 = filtered / total。
        - filtered_types_delete_ratio: 仅在 FILTER_TYPES 子集上的等效删除比例:
              sum(filtered_qtype for qtype in FILTER_TYPES)
            / sum(total_qtype for qtype in FILTER_TYPES)
    """
    if not input_path.is_file():
        raise FileNotFoundError(f"输入 JSON 文件不存在: {input_path}")

    random.seed(seed)

    # 1) 读取整个 JSON (list[dict])
    try:
        raw = json.loads(input_path.read_text(encoding="utf-8"))
    except Exception as e:
        raise RuntimeError(
            f"读取或解析输入 JSON 失败: path={input_path}, error={e}"
        ) from e

    if not isinstance(raw, list):
        raise RuntimeError(
            f"输入 JSON 顶层结构必须为 list, 实际为: {type(raw)} (path={input_path})"
        )

    total = 0
    kept = 0
    filtered = 0

    type_stats: Dict[str, Dict[str, int]] = {}
    output_records: List[Dict[str, Any]] = []

    # 2) 遍历列表中的每条记录
    for idx, record in enumerate(raw, 1):
        if not isinstance(record, dict):
            raise RuntimeError(
                f"索引 {idx} 的元素不是 dict 类型, 实际为: {type(record)}"
            )

        total += 1

        q_type = record.get("question_type", None)
        if not isinstance(q_type, str):
            raise RuntimeError(
                f"索引 {idx} 缺少 question_type 字段或类型非 str: {q_type!r}"
            )

        if q_type not in type_stats:
            type_stats[q_type] = {"total": 0, "kept": 0, "filtered": 0}
        type_stats[q_type]["total"] += 1

        # 非过滤类型 => 直接保留
        if q_type not in FILTER_TYPES:
            output_records.append(record)
            kept += 1
            type_stats[q_type]["kept"] += 1
            continue

        # 过滤类型 => 需要读取图像并根据 keep_ratio 随机采样
        images = record.get("images", None)
        if not isinstance(images, list) or len(images) == 0:
            raise RuntimeError(
                f"索引 {idx} 的 images 字段缺失或为空, 无法获取图像路径: {images!r}"
            )

        img_path = images[0]
        if not isinstance(img_path, str):
            raise RuntimeError(
                f"索引 {idx} 的 images[0] 不是字符串, 实际为: {img_path!r}"
            )
        img_name = os.path.basename(img_path)

        if img_name not in image_keep_ratio:
            raise RuntimeError(
                f"索引 {idx} 的图像文件名在 image_keep_ratio 中未找到: img_name={img_name!r}。"
                f"该图像可能在 support_cleaned 中缺失, 或在“疾病频率模式”下 label 为空/为 None, "
                f"但仍参与 FILTER_TYPES={FILTER_TYPES} 类型的过滤; 按设计不应出现此情况, "
                f"请检查数据准备与过滤逻辑。"
            )

        keep_ratio = image_keep_ratio[img_name]
        r = random.random()
        if r <= keep_ratio:
            output_records.append(record)
            kept += 1
            type_stats[q_type]["kept"] += 1
        else:
            filtered += 1
            type_stats[q_type]["filtered"] += 1

    # 3) 写出过滤后的 JSON (仍然是 list[dict])
    try:
        output_path.write_text(
            json.dumps(output_records, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
    except Exception as e:
        raise RuntimeError(
            f"写出过滤结果 JSON 失败: path={output_path}, error={e}"
        ) from e

    # 基础统计日志
    logger.info(
        "过滤完成: 总样本数=%d, 保留=%d, 丢弃=%d",
        total,
        kept,
        filtered,
    )
    for q_type, st in sorted(type_stats.items(), key=lambda x: x[0]):
        logger.info(
            "  question_type=%s: total=%d, kept=%d, filtered=%d",
            q_type,
            st["total"],
            st["kept"],
            st["filtered"],
        )

    # 节流: 等效删除比例
    if total > 0:
        global_delete_ratio = filtered / float(total)
    else:
        global_delete_ratio = 0.0

    sub_total = 0
    sub_filtered = 0
    for q_type, st in type_stats.items():
        if q_type in FILTER_TYPES:
            sub_total += st["total"]
            sub_filtered += st["filtered"]

    if sub_total > 0:
        sub_delete_ratio = sub_filtered / float(sub_total)
    else:
        sub_delete_ratio = 0.0

    logger.info(
        "等效删除比例(全局): delete_ratio=%.4f (filtered=%d / total=%d)",
        global_delete_ratio,
        filtered,
        total,
    )
    logger.info(
        "等效删除比例(仅 FILTER_TYPES=%s): delete_ratio=%.4f (filtered=%d / total=%d)",
        ",".join(FILTER_TYPES),
        sub_delete_ratio,
        sub_filtered,
        sub_total,
    )

    # 返回统计信息, 供上层生成日志等使用
    stats: Dict[str, Any] = {
        "total": total,
        "kept": kept,
        "filtered": filtered,
        "global_delete_ratio": global_delete_ratio,
        "filtered_types_delete_ratio": sub_delete_ratio,
        "type_stats": type_stats,
    }
    return stats


# -----------------------------
# CLI 入口
# -----------------------------

def parse_args() -> argparse.Namespace:
    """
    命令行参数解析。

    参数:
        --input-jsonl, -i          输入 Alpaca JSON 文件路径 (顶层 list[dict])
        --output-jsonl, -o         输出过滤后的 JSON 文件路径 (顶层 list[dict])
        --support-cleaned-path     支持 JSON 文件路径 (用于疾病统计/固定比例映射)
        --fixed-ratio              固定保留比例 (0,1], 若设置则启用“固定比例模式”
        --seed                     随机种子 (默认 42)
    """
    parser = argparse.ArgumentParser(
        description="基于疾病稀有度或固定比例的 Alpaca 数据过滤工具 (分段线性频率策略 + 可选统一保留率)",
    )
    parser.add_argument(
        "--input-jsonl",
        "-i",
        type=str,
        required=True,
        help="输入 Alpaca JSON 文件路径 (顶层 list[dict])",
    )
    parser.add_argument(
        "--output-jsonl",
        "-o",
        type=str,
        required=True,
        help="输出过滤后的 JSON 文件路径 (顶层 list[dict])",
    )
    parser.add_argument(
        "--support-cleaned-path",
        "-s",
        type=str,
        required=True,
        help="支持 JSON 文件路径, 用于疾病统计或统一 fixed_ratio 映射 (key 为图像文件名)",
    )
    parser.add_argument(
        "--fixed-ratio",
        type=float,
        default=None,
        help="可选: 固定保留比例 p ∈ (0,1]。若指定则启用“固定比例模式”(所有 support 中的图像均采用同一 keep_ratio=p)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="随机种子 (默认 42)",
    )
    return parser.parse_args()


def main() -> None:
    """主函数: 串联支持信息读取、疾病频率/固定比例映射与 Alpaca JSON 过滤。"""
    args = parse_args()

    input_path = Path(args.input_jsonl)
    output_path = Path(args.output_jsonl)
    support_path = Path(args.support_cleaned_path)
    seed = int(args.seed)
    fixed_ratio = args.fixed_ratio

    # 1. 读取支持 JSON
    support_dict = load_support_cleaned(support_path)

    # 2. 根据是否指定 fixed_ratio 决定采用哪种模式构建 image_keep_ratio
    disease_freq: Dict[str, int] = {}
    disease_keep_ratio: Dict[str, float] = {}
    using_adaptive = fixed_ratio is None

    if using_adaptive:
        # ---- 自适应疾病频率模式 ----
        image_to_disease_key, disease_freq = build_disease_key_and_freq(support_dict)
        disease_keep_ratio = build_disease_keep_ratio(disease_freq)
        image_keep_ratio = build_image_keep_ratio(image_to_disease_key, disease_keep_ratio)
        logger.info("当前模式: 基于疾病频率的自适应抽样 (分段线性 + Top3 特殊处理)")
    else:
        # ---- 固定比例模式 ----
        image_keep_ratio = build_image_keep_ratio_fixed(
            support_dict=support_dict,
            fixed_ratio=float(fixed_ratio),
        )
        logger.info("当前模式: 固定比例抽样 fixed_ratio=%.3f (与疾病标签无关)", float(fixed_ratio))

    # 3. 按 keep_ratio 与题目类型过滤 Alpaca JSON (并计算等效删除比例)
    stats = filter_alpaca_json(
        input_path=input_path,
        output_path=output_path,
        image_keep_ratio=image_keep_ratio,
        seed=seed,
    )

    # 4. 若使用“疾病频率模式”, 在输出目录(out_dir)写出额外的 JSON 日志
    #    格式:
    #        {
    #          "等效删除比": <全局等效删除比例>,
    #          "疾病情况": [
    #            [疾病名称, 频次, 删除比例],
    #            ...
    #          ]
    #        }
    #    其中删除比例 = 1 - keep_ratio, 按频次从高到低排序。
    if using_adaptive:
        out_dir = output_path.parent
        # 按频次从高到低排序
        disease_items_sorted = sorted(
            disease_freq.items(),
            key=lambda x: x[1],
            reverse=True,
        )
        disease_list: List[List[Any]] = []
        for d_key, freq in disease_items_sorted:
            kr = disease_keep_ratio.get(d_key)
            if kr is None:
                raise RuntimeError(
                    f"疾病 key={d_key!r} 在 disease_keep_ratio 中未找到, "
                    f"请检查 build_disease_keep_ratio 逻辑"
                )
            delete_ratio = 1.0 - float(kr)
            disease_list.append([d_key, int(freq), float(delete_ratio)])

        log_obj: Dict[str, Any] = {
            "等效删除比": float(stats.get("global_delete_ratio", 0.0)),
            "疾病情况": disease_list,
        }

        log_path = out_dir / f"{output_path.stem}_filter_log.json"
        try:
            out_dir.mkdir(parents=True, exist_ok=True)
            log_path.write_text(
                json.dumps(log_obj, ensure_ascii=False, indent=2),
                encoding="utf-8",
            )
            logger.info("自适应模式下的删除统计日志已写出: %s", log_path)
        except Exception as e:
            raise RuntimeError(
                f"写出过滤日志 JSON 失败: path={log_path}, error={e}"
            ) from e


if __name__ == "__main__":
    main()
